"""
Case Type   : DiskANN向量检索
Case Name   : diskann索引查询过程进行DML、vacuum、analyze等操作，功能正常
Create At   : 2025-09-19
Owner       : taoying
Description :
    1.创建表格并插入数据、构建索引
    2.循环进行vacuum/update操作
    3.循环进行索引查询
Expect      :
    1.成功
    2.成功
    3.成功
History     :
"""
import os
import unittest

from yat.test import Node
from yat.test import macro

from testcase.utils.Common import Common
from testcase.utils.CommonSH import CommonSH
from testcase.utils.CommonSQL import CommonSQL
from testcase.utils.Logger import Logger
from testcase.utils.ComThread import ComThread

logger = Logger()


class DiskANNSearch(unittest.TestCase):
    def setUp(self):
        logger.info('----------------this is setup-----------------------')
        logger.info(
            '---Opengauss_Function_DiskANN_Index_Case0037开始执行-----')
        self.userNode = Node('dbuser')
        self.DB_ENV_PATH = macro.DB_ENV_PATH
        self.comsh = CommonSH('dbuser')
        self.com = Common()
        self.comsql = CommonSQL('dbuser')
        self.table_name = 't_func_diskann_index0037'
        self.fun_insert = 'f_func_diskann_index0037'

    def test_diskann(self):
        step_txt = '----step1:创建表格和diskann索引; expect:成功----'
        sql = f'''
            drop table if exists {self.table_name};
            create table {self.table_name} (id bigserial primary key, val1 vector(5));
            alter table {self.table_name} set (parallel_workers=10);
            drop index if exists idx_t_func_diskann_index0037;
            create index idx_t_func_diskann_index0037 on {self.table_name} using diskann(val1 vector_l2_ops) with(index_size=16);
            '''
        sql_cmd = self.comsh.execut_db_sql(sql)
        logger.info(sql_cmd)

        step_txt = '----step2:创建自定义插入函数; expect:成功----'
        logger.info(step_txt)
        result = self.comsql.create_vector_insert_func(self.fun_insert)
        logger.info(result)
        self.assertIn(self.fun_insert, result, '执行失败:' + step_txt)

        step_txt = '----step3:通过自定义插入函数插入数据; expect:成功----'
        insert_sql = f"select {self.fun_insert}(5, 100, '{self.table_name}', 'val1', 'vector');"
        result = self.comsh.execut_db_sql(insert_sql)
        logger.info(result)
        self.assertIn('SUCCESS', result, '执行失败: 插入数据失败')

        step_txt = '----step4:循环进行vacuum/update操作; expect:成功----'
        excute_cmd1 = '''source {DB_ENV_PATH};
        for i in $(seq 1 10);
        do gsql -d test_db \\
        -p {db_port} \\
        -U {ssh_user} \\
        -c "delete from t_func_diskann_index0037 where id<10; vacuum full t_func_diskann_index0037; update t_func_diskann_index0037 set val1=val1-'[1,1,1,1,1]' where id<10;"
        done;
        '''.format(
            DB_ENV_PATH=self.DB_ENV_PATH,
            db_port=self.userNode.db_port,
            ssh_user=self.userNode.ssh_user
        )
        logger.info(excute_cmd1)
        msg1 = self.userNode.sh(excute_cmd1).result()
        logger.info(msg1)
        self.assertNotIn('Error', msg1)

        step_txt = '----step5:循环进行索引查询; expect:成功----'
        excute_cmd2 = f'''source {self.DB_ENV_PATH};
        for i in $(seq 1 10);
        do gsql -d test_db \\
        -p {self.userNode.db_port} \\
        -U {self.userNode.ssh_user} \\
        -c "set enable_seqscan=off;set enable_indexscan=on;select * from t_func_diskann_index0037 order by val1 <-> '[1,2,3,4,5]' limit 5;"
        done;
        '''
        logger.info(excute_cmd2)
        msg2 = self.userNode.sh(excute_cmd2).result()
        logger.info(msg2)
        self.assertNotIn('Error', msg2)

    def tearDown(self):
        logger.info('----------------this is tearDown-----------------------')
        sql_cmd = self.comsh.execut_db_sql('''drop table t_func_diskann_index0037;''')
        logger.info(sql_cmd)
        logger.info(
            '----Opengauss_Function_DiskANN_Index_Case0037执行完成-----')
